import queue

import cv2
import torch
from yolop_model.yolop import get_net
from torchvision.transforms import transforms
from PIL import ImageGrab
import time
import numpy as np
from multiprocessing import Process,Queue
from loaddata import letterbox_for_img
from control import forward,turn_right,turn_left,stop
normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
    )

transform = transforms.Compose([
            transforms.ToTensor(),
            normalize,
        ])
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
half = device.type != 'cpu'

# origin shaoe 1080x1920
h, w = 1080, 1920
# scale down to 1080/3 and 1920 /3
height, width = 384,640
#
pad_h, pad_w = 12,0
ratio = 0.3333333333333333

def load_net():
    model = get_net()
    checkpoint = torch.load('weights/yolop.pth', map_location=device)
    model.load_state_dict(checkpoint['state_dict'])
    # model = model.cuda()
    model.eval()
    half = device.type != 'cpu'
    if half:
        model.half()  # to FP16
    return model

def read_img(half):
    win = cv2.cvtColor(np.array(ImageGrab.grab(bbox=(0, 26, 800, 626))), cv2.COLOR_BGR2RGB)
    new_win = win[:450, ...]
    new_win = cv2.resize(new_win, (1920, 1080))
    img, ratio, pad = letterbox_for_img(new_win, new_shape=(640,640))
    img = np.ascontiguousarray(img)
    img_tsr = transform(img)
    img_tsr = img_tsr.half() if half else img_tsr.float()  # uint8 to fp16/32
    if img_tsr.ndimension() == 3:
        img_tsr = img_tsr.unsqueeze(0)

    # win = cv2.resize(win,(1920,1080))
    # pic = np.vstack((win[::4,::4],new_win[::4,::4]))
    # cv2.imshow('win', pic)
    # cv2.waitKey(25)
    return img_tsr,new_win

def read_func(queue):
    half = device.type != 'cpu'
    while True:
        imgs = read_img(half)
        time.sleep(0.25)
        if queue.full():
            print('queue is full!')
            # break
        else:
            queue.put(imgs)
        # time.sleep(0.25)

def predict_img(img,model):
    img_tsr, new_win = img
    det_out, da_seg_out, ll_seg_out = model(img_tsr.to(device))
    # print(len(det_out),da_seg_out.size(),ll_seg_out.size())
    da_predict = da_seg_out[:, :, pad_h:(height - pad_h), pad_w:(width - pad_w)]
    # print(da_predict.size())
    da_seg_mask = torch.nn.functional.interpolate(da_predict, scale_factor=int(1 / ratio), mode='bilinear')
    _, da_seg_mask = torch.max(da_seg_mask, 1)
    da_seg_mask = da_seg_mask.int().squeeze().cpu().numpy()
    # da_seg_mask = morphological_process(da_seg_mask, kernel_size=7)
    # print(da_seg_mask.shape)
    ll_predict = ll_seg_out[:, :, pad_h:(height - pad_h), pad_w:(width - pad_w)]
    # print(ll_predict.size())
    ll_seg_mask = torch.nn.functional.interpolate(ll_predict, scale_factor=int(1 / ratio), mode='bilinear')
    _, ll_seg_mask = torch.max(ll_seg_mask, 1)
    ll_seg_mask = ll_seg_mask.int().squeeze().cpu().numpy()
    # print(ll_seg_mask.shape)
    return (da_seg_mask,ll_seg_mask,new_win)

def predect_func(image_queue,model,predection_queue):
    model.to(device)
    while True:
        img_tsr = image_queue.get()
        # print('getting {}'.format(type(img_tsr)))
        # img = img_tsr.numpy()[0]
        # img = np.array(img, dtype=np.uint8)
        # img = np.transpose(img,[1,2,0])
        # print(img.shape)
        # cv2.imshow('win',img)
        # cv2.waitKey(1)
        result = predict_img(img_tsr,model)
        predection_queue.put(result)

def compute_area_square(da_seg_mask):
    pts1 = np.asarray([[4 * 1280 // 9, 3 * 720 // 4], [5 * 1280 // 9, 3 * 720 // 4], [3 * 1280 // 5, 4 * 720 // 5],
                       [2 * 1280 // 5, 4 * 720 // 5]])
    pts2 = np.asarray([[x - 1280 // 6, y] for (x, y) in pts1])
    pts3 = np.asarray([[x + 1280 // 6, y] for (x, y) in pts1])
    pts4 = np.asarray([[x, y - 720 // 15] for (x, y) in pts1])

    s = []
    segment_list = [pts1,pts2,pts3,pts4]
    for pts in segment_list:
        l1 = pts[1][0]-pts[0][0]
        l2 = pts[2][0]-pts[3][0]
        h = pts[3][1]-pts[0][1]
        s.append((l1+l2)/2*h)
    # print(s)
    # the ratio of stop left right forword
    a = []
    for i in range(4):
        mask = np.zeros(da_seg_mask.shape[:2], dtype=np.uint8)
        cv2.fillPoly(mask, np.array([segment_list[i]], dtype=np.int32), (255,))
        res1 = cv2.bitwise_and(da_seg_mask, da_seg_mask, mask=mask)
        a.append(np.sum(res1==255)/s[i])
    return a

def direction_choose(result):
    # print('direction')
    da_seg_mask, ll_seg_mask,img = result
    # print(ll_seg_mask.shape)
    da_seg_mask = np.array(da_seg_mask * 255, dtype=np.uint8)
    # print(da_seg_mask.shape)
    # ll_seg_mask = np.array(ll_seg_mask * 255, dtype=np.uint8)
    da_seg_mask = cv2.resize(da_seg_mask,(1280,720))
    # ll_seg_mask  =cv2.resize(ll_seg_mask,(1280,720))

    pic = cv2.cvtColor(da_seg_mask[::2, ::2],cv2.COLOR_GRAY2BGR)
    # print(img.shape,pic.shape)
    img = cv2.resize(img,(640,360))
    # print(img.shape,pic.shape)
    a = compute_area_square(da_seg_mask)

    # print(pic.shape)
    pic = np.hstack((img,pic))

    direction = 'none'
    if a[3] >= 0.6:
        cv2.putText(pic, 'forward', (200, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 0),3)
        direction = 'forward'

    if a[3]< 0.6 and a[0] >= 0.5:
        cv2.putText(pic, 'forward', (200, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 0), 3)
        direction = 'forward'

    if a[3] < 0.6 and a[0] < 0.5 and a[1]>0.5:
        cv2.putText(pic, 'left', (30, 180), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255),3)
        direction = 'turn_left'
    if a[3] < 0.6 and a[0] < 0.5 and a[2]>0.5:
        cv2.putText(pic, 'right', (480, 180), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 255),3)
        direction = 'turn_right'

    if a[3] < 0.6 and a[0] < 0.5 and a[1]<=0.5:
        cv2.putText(pic, 'stop', (230, 300), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 128, 0),3)
        direction = 'stop'

    if a[3] < 0.6 and a[0] < 0.5 and a[2]<=0.5:
        cv2.putText(pic, 'stop', (230, 300), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 128, 0),3)
        direction = 'stop'

    cv2.imshow('pic', pic)
    cv2.waitKey(1)
    print(pic.shape)
    return direction

def direction_choose_func(predection_queue,direction_queue):
    while True:
        result = predection_queue.get()
        direction = direction_choose(result)
        if direction_queue.full():
            print('queue is full!')
            # break
        else:
            direction_queue.put(direction)
        # time.sleep(0.25)
def direction_execute_func(direction_queue):
    while True:
        func = direction_queue.get()
        if func != 'none':
            print(f'execute {func}')
            eval(func)()
if __name__=='__main__':
    # load model
    # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    model = load_net()
    print('begin')
    image_queue = Queue(maxsize=50)
    predection_queue = Queue(maxsize=100)
    direction_queue = Queue(maxsize=50)

    process_imageReader = Process(target=read_func,args=(image_queue,))
    process_imagePredict = Process(target=predect_func,args=(image_queue,model,predection_queue,))
    process_direction_choose = Process(target=direction_choose_func,args=(predection_queue,direction_queue,))
    process_direction_execute = Process(target=direction_execute_func,args=(direction_queue,))

    process_direction_execute.start()
    process_direction_choose.start()
    process_imagePredict.start()
    process_imageReader.start()

    process_direction_execute.join()
    process_direction_choose.join()
    process_imagePredict.join()
    process_imageReader.join()
    pass
